import sys
import os

# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import rosbag
import numpy as np
import sensor_msgs.point_cloud2 as pc2
import open3d as o3d
import cv2
import math
from cv_bridge import CvBridge
from tictoc import TicToc
import csv
import json
import argparse
import multiprocessing
from utils import topic2path, get_ms
from my_io import write_points, copy_file, files_filter

CONSTANTS_RADIUS_OF_EARTH = 6371000

class Ins2Loc:
    def __init__(self):
        self.is_init = False
        self.latitude = None
        self.longitude = None
        self.altitude = None
        self.vel_x = None
        self.vel_y = None
        self.vel_z = None
        self.accel_x = None
        self.accel_y = None
        self.accel_z = None
        self.roll = None
        self.pitch = None
        self.yaw = None
        self.angular_x = None
        self.angular_y = None
        self.angular_z = None
        self.mercator_x = None
        self.mercator_y = None
        self.mercator_z = None

    def global_init(self, msg):
        self.latitude_init = math.radians(msg.ins_lat)
        self.longitude_init = math.radians(msg.ins_lon)
        self.altitude_init = msg.ins_alt
        self.sin_lat = math.sin(self.latitude_init)
        self.cos_lat = math.cos(self.latitude_init)
        self.is_init = True

    def get_mercator(self, msg):
        latitude = math.radians(msg.ins_lat)
        longitude = math.radians(msg.ins_lon)
        sin_lat = math.sin(latitude)
        cos_lat = math.cos(latitude)
        cos_d_lon = math.cos(longitude - self.longitude_init)

        arg = self.sin_lat * sin_lat + self.cos_lat * cos_lat * cos_d_lon
        if arg > 1.0:
            arg = 1.0
        elif arg < -1.0:
            arg = -1.0

        c = math.acos(arg)
        if math.fabs(c) < 0.0001:
            k = 1.0
        else:
            k = c / math.sin(c)

        self.mercator_x = k * cos_lat * math.sin(longitude - self.longitude_init) * CONSTANTS_RADIUS_OF_EARTH
        self.mercator_y = k * (self.cos_lat * sin_lat - self.sin_lat * cos_lat * cos_d_lon) * CONSTANTS_RADIUS_OF_EARTH
        self.mercator_z = self.altitude_init - msg.ins_alt

    def forward(self, msg):
        self.latitude = msg.ins_lat
        self.longitude = msg.ins_lon
        self.altitude = msg.ins_alt
        self.vel_x = msg.ins_vel_e
        self.vel_y = msg.ins_vel_n
        self.vel_z = msg.ins_vel_u
        self.accel_x = msg.acc_x
        self.accel_y = msg.acc_y
        self.accel_z = msg.acc_z
        self.roll = msg.ins_roll
        self.pitch = msg.ins_pitch
        self.yaw = msg.ins_yaw
        self.angular_x = msg.gyro_x
        self.angular_y = msg.gyro_y
        self.angular_z = msg.gyro_z
        if not self.is_init:
            self.global_init(msg)
        self.get_mercator(msg)




class MapProjectionReference:
    """参考点初始化类（类似 C 的 struct map_projection_reference_s）"""

    def __init__(self):
        self.lat_rad = 0.0
        self.lon_rad = 0.0
        self.sin_lat = 0.0
        self.cos_lat = 0.0
        self.alt = 0.0
        self.yaw = 0.0
        self.mercator_x = 0.0
        self.mercator_y = 0.0
        self.mercator_z = 0.0
        self.init_flag = False

    def is_initialized(self):
        """检查参考点是否已初始化"""
        return self.init_flag


def map_projection_init(ref, lat_deg, lon_deg, alt_0):
    """
    初始化参考点
    :param ref: MapProjectionReference 对象
    :param lat_deg: 参考点纬度（度）
    :param lon_deg: 参考点经度（度）
    """
    ref.lat_rad = math.radians(lat_deg)
    ref.lon_rad = math.radians(lon_deg)
    ref.sin_lat = math.sin(ref.lat_rad)
    ref.cos_lat = math.cos(ref.lat_rad)
    ref.alt = alt_0
    ref.init_flag = True


def map_projection_project(ref, lat_deg, lon_deg, alt, yaw):
    """
    将经纬度转换为局部切平面坐标
    :param ref: 已初始化的 MapProjectionReference 对象
    :param lat_deg: 目标点纬度（度）
    :param lon_deg: 目标点经度（度）
    :return: (x, y) 东向和北向坐标（米），若失败返回 (None, None)
    """
    if not ref.is_initialized():
        return None

    # 将经纬度转换为弧度
    lat_rad = math.radians(lat_deg)
    lon_rad = math.radians(lon_deg)

    # 计算相对经度差
    d_lon = lon_rad - ref.lon_rad
    cos_d_lon = math.cos(d_lon)

    # 计算球面角距
    sin_lat = math.sin(lat_rad)
    cos_lat = math.cos(lat_rad)
    arg = ref.sin_lat * sin_lat + ref.cos_lat * cos_lat * cos_d_lon

    # 约束 arg 在 [-1, 1] 范围内
    arg = max(min(arg, 1.0), -1.0)

    c = math.acos(arg)
    if abs(c) < 1e-15:  # 处理 c ≈ 0 的情况，避免除以零
        k = 1.0
    else:
        k = c / math.sin(c)

    # 计算平面坐标
    x = k * (ref.cos_lat * sin_lat - ref.sin_lat * cos_lat * cos_d_lon) * CONSTANTS_RADIUS_OF_EARTH
    y = k * cos_lat * math.sin(d_lon) * CONSTANTS_RADIUS_OF_EARTH

    local_pt = MapProjectionReference()

    map_projection_init(local_pt, lat_deg, lon_deg, alt)
    local_pt.mercator_x = y
    local_pt.mercator_y = x
    local_pt.mercator_z = ref.alt - alt
    local_pt.yaw = yaw

    return local_pt

def main(args):
    # 统计耗时
    cost = TicToc("ros解析")
    assert os.path.exists(args.rosbag_dir)
    files = files_filter(args.rosbag_dir, pos_fix="bag", split='.')
    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args.rosbag_dir, args.data_path))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args.rosbag_dir, args.data_path)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(bag, rosbag_dir, data_path):
    rosbag_path = os.path.join(rosbag_dir, bag)
    dir_name = os.path.join(data_path, bag.split('.')[0])
    print("开始解析<{}>......".format(rosbag_path))

    bag = rosbag.Bag(rosbag_path)
    #
    # lidar_topics = ["/pointcloud_lidar3"]
    # radar_topics = ["/pointcloud_radar3"]
    # old_radar_topics = {"/radar_det_list": "/pointcloud_radar3"}
    # # camera_topics = ["/camera71/compressed", "/camera72/compressed", "/camera73/compressed", "/camera74/compressed",
    # #                  "/camera75/compressed", "/camera76/compressed",
    # #                  "/camera77/compressed", "/camera78/compressed", "/camera79/compressed", "/camera80/compressed",
    # #                  "/camera81/compressed", "/camera82/compressed"]
    # camera_topics = ["/camera75/compressed"]
    # camera_ori_topics = []

    lidar_topics = []
    radar_topics = []
    old_radar_topics = {}
    # camera_topics = ["/camera71/compressed", "/camera72/compressed", "/camera73/compressed", "/camera74/compressed",
    #                  "/camera75/compressed", "/camera76/compressed",
    #                  "/camera77/compressed", "/camera78/compressed", "/camera79/compressed", "/camera80/compressed",
    #                  "/camera81/compressed", "/camera82/compressed"]
    camera_topics = []
    camera_ori_topics = []


    localization_topics = ["/ins_msg_dy"]
    localer = Ins2Loc()
    print("保存在<{}>".format(dir_name))

    # 新建目录
    new_dir = []
    ffarther_dir = []
    farther_dir = []
    child_dir = []
    ffarther_dir.append(dir_name)
    new_dir = new_dir + ffarther_dir
    farther_dir.append(os.path.join(dir_name, "lidar"))
    farther_dir.append(os.path.join(dir_name, "radar"))  # 毫米波所以直接保存为bin:因为1)不直接标注 2)pcd不能保留r_speed和rcs
    farther_dir.append(os.path.join(dir_name, "camera"))
    farther_dir.append(os.path.join(dir_name, "localization"))
    farther_dir.append(os.path.join(dir_name, "calib"))
    new_dir = new_dir + farther_dir
    for topic in lidar_topics:
        child_dir.append(os.path.join(dir_name, "lidar", topic2path(topic)))
    for topic in radar_topics:
        child_dir.append(os.path.join(dir_name, "radar", topic2path(topic)))
    for topic in camera_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic, -2)))
    for topic in camera_ori_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic)))
    new_dir = new_dir + child_dir
    for dir in new_dir:
        if not os.path.exists(dir):
            print("新建目录<{}>".format(dir))
            os.makedirs(dir, exist_ok=True)

    if os.path.exists(os.path.join(rosbag_dir, 'calib.json')):
        copy_file(os.path.join(rosbag_dir, 'calib.json'), os.path.join(dir_name, 'calib'))
        print("拷贝 calib.json , from {} to {}".format(os.path.join(rosbag_dir, 'calib'),
                                                       os.path.join(dir_name, 'calib')))
    else:
        # 创建calib.yml
        eye44 = np.eye(4).tolist()  # 4*4的旋转平移矩阵
        eye33 = np.eye(3).tolist()  # 3*3的旋转平移矩阵
        calib_file = os.path.join(dir_name, 'calib', 'calib.json')
        calib_data = dict()
        for topic in lidar_topics + radar_topics:
            sensor_id = topic[-6:]
            calib_data[sensor_id] = dict()
            calib_data[sensor_id]['transforms'] = dict()
            calib_data[sensor_id]['transforms']['base_link'] = eye44

        calib_data['camera75'] = dict()
        calib_data['camera75']['ori'] = dict()
        calib_data['camera75']['ori']['imgh'] = 9999
        calib_data['camera75']['ori']['imgw'] = 9999
        calib_data['camera75']['ori']['K'] = eye33
        calib_data['camera75']['ori']['D'] = np.zeros(shape=(0, 5)).tolist()

        # 写入JSON文件
        with open(calib_file, 'w', encoding='utf-8') as file:
            json.dump(calib_data, file, ensure_ascii=False, indent=4)

        # fs = cv2.FileStorage(os.path.join(dir_name, 'calib', 'sensor.yml'), cv2.FileStorage_WRITE)
        # matrix = np.eye(4)  # 4*4的旋转平移矩阵
        # for topic in lidar_topics + radar_topics:
        #     sensor_id = topic[-6:]
        #     fs.write(sensor_id, matrix)
        # fs.release()
        print("新建标定文件 in {}".format(calib_file))

    # 创建或打开文件
    init_pos = MapProjectionReference() #初始位置
    csvfile = open(os.path.join(dir_name, "localization", 'localization.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['stamp(ms)', 'latitude', 'longitude', 'altitude', 'mercator_x', 'mercator_y', 'mercator_z',
                  'vel_x', 'vel_y', 'vel_z', 'accel_x', 'accel_y', 'accel_z',
                  'roll', 'pitch', 'yaw', 'angular_x', 'angular_y', 'angular_z']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()

    for topic, msg, t in bag:
        if topic in lidar_topics:
            lidar = pc2.read_points(msg)
            points = np.array(list(lidar))
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式 比如用于lc标定任务
            # o3d.io.write_point_cloud(os.path.join(dir_name,'lidar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)
            assert points.shape[1] >= 4
            write_points(points[:, [0, 1, 2, 3, 5]],
                         os.path.join(dir_name, 'lidar', topic2path(topic), str(get_ms(msg)) + ".bin"))

        elif topic in radar_topics:
            radar = pc2.read_points(msg)
            points = np.array(list(radar))
            points = points[:, :5]  # x y z r_speed rcs
            file_path = os.path.join(dir_name, 'radar', topic2path(topic), str(get_ms(msg)) + ".bin")
            write_points(points, file_path)
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])# x y z r_speed rcs
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式
            # o3d.io.write_point_cloud(os.path.join(dir_name,'radar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)

        elif topic in old_radar_topics:
            points_list = []
            for det in msg.detection_array:
                points_list.append([det.f_x, det.f_y, det.f_z, det.f_RangeRate, det.s_RCS])
            points = np.array(points_list)
            file_path = os.path.join(dir_name, 'radar', topic2path(old_radar_topics[topic]),
                                     str(get_ms(msg.detection_array[0])) + ".bin")
            write_points(points, file_path)


        elif topic in camera_topics:
            np_arr = np.frombuffer(msg.data, dtype=np.uint8)  # 将消息中的字节数据转换为 numpy 数组
            cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg)+500) + ".jpg"),cv_image)
            # cv2.imwrite(
            #     os.path.join(dir_name, "camera", topic2path(topic, index=-2), str(get_ms(msg)) + ".jpg"),
            #     cv_image)

        # elif topic in camera_topics:
        #     bridge = CvBridge()
        #     cv_image = bridge.compressed_imgmsg_to_cv2(msg, "bgr8")
        #     # tm4
        #     # if topic[:9] == '/camera81' or topic[:9] == '/camera80':  # 前后相机超前600ms，左右相机超前500ms
        #     #     cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 450) + ".jpg"),
        #     #                 cv_image)
        #     # else:
        #     #     cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 600) + ".jpg"),
        #     #                 cv_image)
        #     # kb qj101
        #     if topic[:9] == '/camera81' or topic[:9] == '/camera80':  # 前后相机超前600ms，左右相机超前500ms
        #         cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 650) + ".jpg"),
        #                     cv_image)
        #     else:
        #         cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 750) + ".jpg"),
        #                     cv_image)

            # cv2.imshow("Image window", cv_image)
            # cv2.waitKey(3)
        # elif topic in camera_ori_topics:
        #     bridge = CvBridge()
        #     cv_image = bridge.imgmsg_to_cv2(msg, "bgr8")
        #     # cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg)+600) + ".jpg"), cv_image)
        #     if topic[:9] == '/camera81':  # 前后相机超前600ms，左右相机超前500ms
        #         cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 450) + ".jpg"),
        #                     cv_image)
        #     else:
        #         cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg) + 600) + ".jpg"),
        #                     cv_image)

        elif topic in localization_topics:

            localer.forward(msg=msg)

            # +'\t'用于取消科学计数法
            write.writerow({'stamp(ms)': str(get_ms(msg)) + '\t',
                            'latitude': localer.latitude,
                            'longitude': localer.longitude,
                            'altitude': localer.altitude,
                            'mercator_x': localer.mercator_x,
                            'mercator_y': localer.mercator_y,
                            'mercator_z': localer.mercator_z,
                            'vel_x': localer.vel_x,
                            'vel_y': localer.vel_y,
                            'vel_z': localer.vel_z,
                            'accel_x': localer.accel_x,
                            'accel_y': localer.accel_y,
                            'accel_z': localer.accel_z,
                            'roll': localer.roll,
                            'pitch': localer.pitch,
                            'yaw': localer.yaw,
                            'angular_x': localer.angular_x,
                            'angular_y': localer.angular_y,
                            'angular_z': localer.angular_z})

            # if not init_pos.is_initialized():
            #     map_projection_init(init_pos, msg.ins_lat, msg.ins_lon, msg.ins_alt)
            #     local_pt = init_pos
            # else:
            #     local_pt = map_projection_project(init_pos, msg.ins_lat, msg.ins_lon, msg.ins_alt, msg.ins_yaw)
            #
            # # +'\t'用于取消科学计数法
            # write.writerow({'stamp(ms)': str(get_ms(msg)) + '\t',
            #                 'latitude': msg.ins_lat, 'longitude': msg.ins_lon, 'altitude': msg.ins_alt,
            #                 'mercator_x': local_pt.mercator_x, 'mercator_y': local_pt.mercator_y, 'mercator_z': local_pt.mercator_z,
            #                 'vel_x': msg.ins_vel_e, 'vel_y': msg.ins_vel_n, 'vel_z': msg.ins_vel_u,
            #                 'accel_x': msg.acc_x, 'accel_y': msg.acc_y, 'accel_z': msg.acc_z,
            #                 'roll': msg.ins_roll, 'pitch': msg.ins_pitch, 'yaw': msg.ins_yaw,
            #                 'angular_x': msg.gyro_x, 'angular_y': msg.gyro_y, 'angular_z': msg.gyro_z})

    for dir in child_dir:
        print("文件夹:<{}> 包含文件<{}>个".format(dir.split('/')[-1], len(os.listdir(dir))))

    bag.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--rosbag-dir',
                        default='/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/tmx/bags',
                        help='rosbag')
    parser.add_argument('--data-path',
                        default='/media/adt/ZWH4T/ZWH/Dataset3d/dataset_18xx-bag/cyw2.0/tmx/datasets')
    args = parser.parse_args()

    main(args)
